from dataset.dataset import DRR_dataset
import torch
from utils.logger import Logger
from network.ResNet_model import resnet50
from torch import nn, optim
import network.metrics as met


# 计算loss
def get_loss(pre, tru):
    """
    根据网络的预测值和真实值，输出损失
    :param pre: 预测值
    :param tru: 真实值
    :return: 损失
    """
    pre_rx = pre[:, 0: 6]
    pre_ry = pre[:, 6: 12]
    pre_rz = pre[:, 12: 18]
    pre_tx = pre[:, 18: 30]
    pre_ty = pre[:, 30: 42]
    pre_tz = pre[:, 42: 62]
    # 将真实值转换为类别
    tru_rx = torch.clamp(tru[:, 0:3] / 5 - 1, min=0).long()
    tru_txy = torch.clamp(tru[:, 3:5] / 2.5 - 1, min=0).long()
    tru_tz = torch.clamp(tru[:, 5] / 2 - 1, min=0).long()
    soft_loss = nn.CrossEntropyLoss()
    loss1 = soft_loss(pre_rx, tru_rx[:, 0])
    loss2 = soft_loss(pre_ry, tru_rx[:, 1])
    loss3 = soft_loss(pre_rz, tru_rx[:, 2])
    loss4 = soft_loss(pre_tx, tru_txy[:, 0])
    loss5 = soft_loss(pre_ty, tru_txy[:, 1])
    loss6 = soft_loss(pre_tz, tru_tz)
    # 加上权重系数，使其更加注重角度上的偏移量
    loss = 1.1 * (loss1 + loss2 + loss3) + loss4 + loss5 + loss6
    # 将类别转换为角度和位移输出
    softmax = nn.Softmax(dim=1)
    _, out_rx = torch.max(softmax(pre_rx), dim=1)
    # print(out_rx.unsqueeze(dim=0).mT)
    _, out_ry = torch.max(softmax(pre_ry), dim=1)
    _, out_rz = torch.max(softmax(pre_rz), dim=1)
    _, out_tx = torch.max(softmax(pre_tx), dim=1)
    _, out_ty = torch.max(softmax(pre_ty), dim=1)
    _, out_tz = torch.max(softmax(pre_tz), dim=1)
    out_r = torch.cat((out_rx.unsqueeze(dim=0).mT, out_ry.unsqueeze(dim=0).mT, out_rz.unsqueeze(dim=0).mT), dim=1) * 5 + 2.5
    out_txy = torch.cat((out_tx.unsqueeze(dim=0).mT, out_ty.unsqueeze(dim=0).mT), dim=1) * 2.5 + 1.25
    out = torch.cat((out_r, out_txy, out_tz.unsqueeze(dim=0).mT * 2 + 1), dim=1)
    return loss, out


# 真正实际运行的主程序
class ClaTrain:
    def __init__(self, num_classes=62, batch_size=8,
                 train_path=None, train_label_path=None,
                 test_path=None, test_label_path=None,
                 ):
        # 数据导入
        self.train_dataset = DRR_dataset(DRR_path=train_path, fn_path=train_label_path,
                                         Batch_size=batch_size)
        self.test_dataset = DRR_dataset(DRR_path=test_path, fn_path=test_label_path,
                                        Batch_size=batch_size)
        # 训练模型
        self.train_model = resnet50(num_classes=num_classes)
        # 定义损失函数和优化器
        self.optimizer = optim.Adam(self.train_model.parameters(), lr=0.0001)
        # 训练的相关参数
        self.init_epoch = 0  # 轮次
        self.Batch_size = batch_size
        self.min_loss = 10000  # 设置最小损失，用于判断并保存断点

    def train_once(self):
        # 前向传播一轮
        self.train_model.train().cuda()
        batch_x, batch_y = self.train_dataset.get_data()
        out = self.train_model(batch_x)
        loss, out = get_loss(out, batch_y)
        # 反向传播更新权重参数
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return batch_y, out, loss

    def test(self):
        # 在测试集上测试
        self.train_model.eval()
        with torch.no_grad():
            batch_x, batch_y = self.test_dataset.get_data()
            out = self.train_model(batch_x)
            loss, out = get_loss(out, batch_y)
        return batch_y, out, loss


if __name__ == "__main__":
    pass
